import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import sigmaclip

class BiasImageAnalyzer:
    def __init__(self, image_data, grid_size=(16, 16), clip_sigma=None):
        """
        初始化图像分析器

        参数:
        - image_data: 2D numpy array，bias 或 dark 图像数据
        - grid_size: tuple(int, int)，用于分块的网格数，如 (16, 16)
        - clip_sigma: float 或 None，是否启用 sigma clipping，若为 None 则不启用
        """
        self.grid_y, self.grid_x = grid_size
        self.image_data = image_data
        self.clip_sigma = clip_sigma
        self.mean_map = None
        self.std_map = None

    def sigma_clip(self, data):
        """
        对一维数据进行 sigma clipping

        参数:
        - data: 一维 numpy array

        返回:
        - clipped: 去除异常值后的数据
        """
        if self.clip_sigma is not None:
            clipped, _, _ = sigmaclip(data, low=self.clip_sigma, high=self.clip_sigma)
            return clipped
        else:
            return data

    def compute_mean_std_maps_mean(self):
        """
        将图像按网格划分，并计算每个网格的均值和标准差（可选 sigma clipping）
        """
        H, W = self.image_data.shape
        block_h = H // self.grid_y
        block_w = W // self.grid_x
        self.mean_map = np.zeros((self.grid_y, self.grid_x))
        self.std_map = np.zeros((self.grid_y, self.grid_x))
        for i in range(self.grid_y):
            for j in range(self.grid_x):
                block = self.image_data[i*block_h:(i+1)*block_h, j*block_w:(j+1)*block_w]
                clipped_block = self.sigma_clip(block.ravel())
                self.mean_map[i, j] = np.mean(clipped_block)
                self.std_map[i, j] = np.std(clipped_block)
    def compute_mean_std_maps(self):
        """
        将图像按网格划分，并计算每个网格的中值和标准差（可选 sigma clipping）
        """
        H, W = self.image_data.shape
        block_h = H // self.grid_y
        block_w = W // self.grid_x
        self.mean_map = np.zeros((self.grid_y, self.grid_x))
        self.std_map = np.zeros((self.grid_y, self.grid_x))
        for i in range(self.grid_y):
            for j in range(self.grid_x):
                block = self.image_data[i*block_h:(i+1)*block_h, j*block_w:(j+1)*block_w]
                clipped_block = self.sigma_clip(block.ravel())
                self.mean_map[i, j] = np.median(clipped_block)  # ← 使用中值
                self.std_map[i, j] = np.std(clipped_block)
                
    def plot_mean_map(self, title=None):
        """
        绘制均值热图

        参数:
        - title: 图标题，默认自动根据 clip_sigma 设置
        """
        if title is None:
            title = f"Mean Map ({self.clip_sigma}σ clipped)" if self.clip_sigma else "Mean Map"
        plt.figure(figsize=(8, 6))
        plt.imshow(self.mean_map, cmap='viridis', origin='lower')
        plt.colorbar(label='Mean ADU')
        plt.title(title)
        plt.xlabel('Grid X')
        plt.ylabel('Grid Y')
        for (j, i), label in np.ndenumerate(self.mean_map):
            plt.text(i, j, f"{label:.1f}", ha='center', va='center', fontsize=6, color='white')
        plt.tight_layout()
        plt.show()

    def plot_std_map(self, title=None):
        """
        绘制标准差热图

        参数:
        - title: 图标题，默认自动根据 clip_sigma 设置
        """
        if title is None:
            title = f"Standard Deviation Map ({self.clip_sigma}σ clipped)" if self.clip_sigma else "Standard Deviation Map"
        plt.figure(figsize=(8, 6))
        plt.imshow(self.std_map, cmap='magma', origin='lower')
        plt.colorbar(label='Standard Deviation')
        plt.title(title)
        plt.xlabel('Grid X')
        plt.ylabel('Grid Y')
        plt.tight_layout()
        plt.show()

    def plot_profile(self):
        """
        绘制行和列方向的均值剖面图，用于分析是否存在系统性偏差或结构
        """
        mean_row = np.mean(self.image_data, axis=1)
        mean_col = np.mean(self.image_data, axis=0)
        plt.figure(figsize=(10, 4))

        plt.subplot(1, 2, 1)
        plt.plot(mean_row)
        plt.title('Mean along Rows')
        plt.xlabel('Row')
        plt.ylabel('Mean ADU')

        plt.subplot(1, 2, 2)
        plt.plot(mean_col)
        plt.title('Mean along Columns')
        plt.xlabel('Column')
        plt.ylabel('Mean ADU')

        plt.tight_layout()
        plt.show()
